"""
手写体识别
"""
import numpy as np
import paddle
import matplotlib.pyplot as plt
import paddle.nn.functional as F

# 确保从paddle.vision.datasets.MNIST中加载的图像数据是np.ndarray类型
paddle.vision.set_image_backend('cv2')

# train_dataset = paddle.vision.datasets.MNIST(mode='train')
#
# train_data_0 = np.array(train_dataset[0][0])
# train_label_0 = np.array(train_dataset[0][1])
#
# plt.figure('Image')
# # plt.figure(figsize=(2,2))
# plt.imshow(train_data_0,cmap='gray')
# plt.axis('on')
# plt.title('image')
#
# # plt.show()
#
# print("图像数据形状和对应数据为：",train_data_0.shape)
# print("图像标签形状和对应数据为：",train_label_0.shape,train_label_0)
# print("\n打印第一个batch的第一个图像，对应标签数字为{}".format(train_label_0))

class MNIST(paddle.nn.Layer):
    """
    定义模型对象
    """
    def __init__(self):
        super(MNIST,self).__init__()

        # 定义一层全连接层，输出维度是1
        self.fc = paddle.nn.Linear(in_features=784,out_features=1)

    # 定义网络结构的前向计算过程
    def forward(self,inputs):
        outputs = self.fc(inputs)
        return outputs

def norm_img(img):
    """
    图像归一化处理
    :param img:
    :return:
    """
    assert len(img.shape) == 3
    batch_size,img_h,img_w = img.shape[0],img.shape[1],img.shape[2]
    # 归一化图像数据
    img = paddle.reshape(img,[batch_size,img_h*img_w])
    return img


def train(model):
    # 启动训练模式
    model.train()
    # 加载训练集 batch_size 设为100
    train_loader = paddle.io.DataLoader(paddle.vision.datasets.MNIST(mode='train'),batch_size=50,shuffle=True)
    # print(enumerate(train_loader()))
    # for batch_id,data in enumerate(train_loader()):
    #     print("batch: {},data:{}".format(batch_id,data))
    # 定义优化器，使用随机梯度下降SGD优化器，学习率设置为0.001
    opt = paddle.optimizer.SGD(learning_rate=0.001,parameters=model.parameters())
    epoch_num = 10
    loss_list = []
    print("开始循环")
    for epoch in range(epoch_num):
        for batch_id,data in enumerate(train_loader()):
            # print("data:",data)
            images = np.array([norm_img(data[0]).astype('float32')])
            labels = np.array([data[1].astype('float32')])

            # images, labels = data
            images = paddle.to_tensor(images)
            labels = paddle.to_tensor(labels)

            # 前向计算的过程
            predicts = model(images)

            # 计算损失
            loss = F.square_error_cost(predicts,labels)
            # print(loss)
            avg_loss = paddle.mean(loss)

            # 没训练了1000批次的数据，打印下当前Loss的情况
            if batch_id % 1000 == 0:
                loss = avg_loss.numpy()
                loss_list.append(loss)
                print("epoch_id: {},batch_id: {},loss is: {}".format(epoch,batch_id,loss))

            # 后向传播，更新参数的过程
            avg_loss.backward()
            opt.step()
            opt.clear_grad()

    return loss_list

model = MNIST()
loss_list = train(model)
# paddle.save(model.state_dict(),'./mnist.pdparams')